import argparse

import mmcv
import models  # noqa: F401,F403
import torch
import torch.nn as nn
from mmcv.runner import load_checkpoint
from mmdet.models import build_detector


class GFLONNXModel(nn.Module):
    def __init__(self, cfg_file_path: str, checkpoint_file_path: str):
        super().__init__()
        self.cfg = mmcv.Config.fromfile(cfg_file_path)
        self.model = self.create_model(cfg_file_path, checkpoint_file_path)
        self.mean = torch.tensor(self.cfg.img_norm_cfg.mean)
        self.mean = self.mean[None, :, None, None]  # shape is: 1 x 3 x 1 x 1
        self.std = torch.tensor(self.cfg.img_norm_cfg.std)
        self.std = self.std[None, :, None, None]  # shape is: 1 x 3 x 1 x 1

    def create_model(self, cfg_file_path, checkpoint_file_path):
        self.cfg.model.pretrained = None
        self.cfg.data.test.test_mode = True
        model = build_detector(self.cfg.model,
                               train_cfg=None,
                               test_cfg=self.cfg.test_cfg)
        load_checkpoint(model, checkpoint_file_path, map_location="cpu")
        model.cpu().eval()
        return model

    def forward(self, img):
        img.sub_(self.mean).div_(self.std)
        with torch.no_grad():
            x = self.model.extract_feat(img)
            cls_score_list, bbox_pred_list = self.model.bbox_head(x)
        return cls_score_list, bbox_pred_list


def argument_parser():
    parser = argparse.ArgumentParser(description="export model to onnx")
    parser.add_argument("--config-file",
                        default="",
                        metavar="FILE",
                        help="path to config file")
    parser.add_argument("--checkpoint",
                        default="",
                        metavar="FILE",
                        help="path to checkpoint file")
    parser.add_argument("--onnx-file",
                        default="out.onnx",
                        help="path to onnx file")
    args = parser.parse_args()
    return args


def main():
    args = argument_parser()
    model = GFLONNXModel(args.config_file, args.checkpoint)

    image = mmcv.imread("images/bird.jpg")
    image = mmcv.imresize(image, size=(608, 608))
    image = image.transpose(2, 0, 1)
    image = torch.from_numpy(image).unsqueeze(0).float()

    input_names = ["image"]
    dynamic_axes = {
        "image": {
            2: "height",
            3: "width"
        },
    }
    torch.onnx.export(model,
                      image,
                      args.onnx_file,
                      input_names=input_names,
                      dynamic_axes=dynamic_axes,
                      export_params=True,
                      keep_initializers_as_inputs=True,
                      verbose=True,
                      opset_version=11)


if __name__ == "__main__":
    main()
